import os
import numpy as np
import rasterio
from rasterio.warp import reproject, Resampling
from rasterio.mask import mask
from scipy.stats import pearsonr
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize
import geopandas as gpd
from tqdm import trange
"""
自变量Z年份，Y降雨量或气温，X因变量NPP
"""
def resample_to_match(source_data, source_transform, target_shape, target_transform, source_crs, target_crs):
    """重采样函数：将源数据重采样到目标分辨率"""
    resampled_data = np.empty(target_shape, dtype=np.float32)
    reproject(
        source=source_data,
        destination=resampled_data,
        src_transform=source_transform,
        dst_transform=target_transform,
        src_crs=source_crs,
        dst_crs=target_crs,
        resampling=Resampling.bilinear  # 使用双线性插值
    )
    return resampled_data

def load_tif_stack(path, mask_gdf=None, target_meta=None):
    """加载时间序列 tif 文件，可选进行重采样"""
    files = sorted([os.path.join(path, f) for f in os.listdir(path) if f.endswith('.tif')])
    stack = []
    
    for file in files:
        with rasterio.open(file) as src:
            if mask_gdf is not None:
                out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
                data = out_image[0]
                transform = out_transform
            else:
                data = src.read(1)
                transform = src.transform
            
            # 如果提供了目标元数据，进行重采样
            if target_meta is not None:
                data = resample_to_match(
                    source_data=data,
                    source_transform=transform,
                    target_shape=(target_meta['height'], target_meta['width']),
                    target_transform=target_meta['transform'],
                    source_crs=src.crs,
                    target_crs=target_meta['crs']
                )
                transform = target_meta['transform']
            
            nodata = src.nodata
            if nodata is not None:
                data = np.where(data == nodata, np.nan, data)
            stack.append(data)
            
            # 保存第一个文件的元数据
            if len(stack) == 1:
                crs = src.crs
                meta = src.meta.copy()
                if target_meta is not None:
                    meta.update({
                        'width': target_meta['width'],
                        'height': target_meta['height'],
                        'transform': target_meta['transform']
                    })
    
    return np.stack(stack), transform, crs, meta

def calculate_partial_correlation(npp_data, climate_data, years):
    """计算偏相关系数"""
    # 确保没有 NaN 值
    valid_mask = ~np.isnan(npp_data) & ~np.isnan(climate_data) & ~np.isnan(years)
    if np.sum(valid_mask) < 3:  # 至少需要3个有效值
        return np.nan
    
    npp_valid = npp_data[valid_mask]
    climate_valid = climate_data[valid_mask]
    years_valid = years[valid_mask]
    
    # 计算简单相关系数
    rxy, _ = pearsonr(npp_valid, climate_valid)  # NPP 与气候因子的相关系数
    rxz, _ = pearsonr(npp_valid, years_valid)    # NPP 与年份的相关系数
    ryz, _ = pearsonr(climate_valid, years_valid) # 气候因子与年份的相关系数
    
    # 计算偏相关系数
    numerator = rxy - (rxz * ryz)
    denominator = np.sqrt((1 - rxz**2) * (1 - ryz**2))
    
    if denominator == 0:
        return np.nan
    
    return numerator / denominator

# 文件路径
npp_path = './mean_npp/'           # NPP 文件夹路径
rainfall_path = './mean_rain_tif/'     # 降雨量文件夹路径
temperature_path = './mean_tem_tif/'   # 气温文件夹路径
mask_shp = './shanxi_mask/shanxi.shp'  # 陕西 mask 文件路径
output_path = './correlation_outputs/'  # 输出文件夹路径

# 确保输出文件夹存在
os.makedirs(output_path, exist_ok=True)

# 加载陕西 mask
mask_gdf = gpd.read_file(mask_shp)
# 首先加载 NPP 数据以获取目标分辨率信息
npp_stack, npp_transform, npp_crs, npp_meta = load_tif_stack(npp_path)
# 加载数据
# npp_stack, npp_transform, npp_crs, npp_meta = load_tif_stack(npp_path)
# 加载并重采样降雨量和气温数据到 NPP 的分辨率
rainfall_stack, _, _, _ = load_tif_stack(
    rainfall_path, 
    mask_gdf=mask_gdf, 
    target_meta={
        'width': npp_meta['width'],
        'height': npp_meta['height'],
        'transform': npp_transform,
        'crs': npp_crs
    }
)

temperature_stack, _, _, _ = load_tif_stack(
    temperature_path, 
    mask_gdf=mask_gdf, 
    target_meta={
        'width': npp_meta['width'],
        'height': npp_meta['height'],
        'transform': npp_transform,
        'crs': npp_crs
    }
)

# 生成年份数组
years = np.arange(2002, 2002 + len(npp_stack))

# 初始化结果数组
rows, cols = npp_stack[0].shape
partial_corr_rainfall = np.full((rows, cols), np.nan)
partial_corr_temperature = np.full((rows, cols), np.nan)

# 计算每个像素的偏相关系数
for i in trange(rows):
    for j in range(cols):
        # 提取时间序列
        npp_series = npp_stack[:, i, j]
        rainfall_series = rainfall_stack[:, i, j]
        temperature_series = temperature_stack[:, i, j]
        
        # 计算降雨量的偏相关系数
        partial_corr_rainfall[i, j] = calculate_partial_correlation(
            npp_series, rainfall_series, years)
        
        # 计算气温的偏相关系数
        partial_corr_temperature[i, j] = calculate_partial_correlation(
            npp_series, temperature_series, years)

# 准备保存结果的元数据
meta = {
    'driver': 'GTiff',
    'dtype': 'float32',
    'nodata': np.nan,
    'width': cols,
    'height': rows,
    'count': 1,
    'crs': npp_crs,
    'transform': npp_transform
}

# 保存结果
with rasterio.open(os.path.join(output_path, 'partial_correlation_rainfall.tif'), 'w', **meta) as dst:
    dst.write(partial_corr_rainfall.astype('float32'), 1)

with rasterio.open(os.path.join(output_path, 'partial_correlation_temperature.tif'), 'w', **meta) as dst:
    dst.write(partial_corr_temperature.astype('float32'), 1)

print("偏相关系数计算完成并保存！")

print("偏相关系数计算完成并保存！")

# 添加可视化代码
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 创建一个包含两个子图的图形
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))

# 绘制降雨量偏相关系数图
im1 = ax1.imshow(partial_corr_rainfall, 
                 cmap='RdBu_r',  # 使用红蓝色图，_r表示反转颜色
                 vmin=-1, 
                 vmax=1)
ax1.set_title('NPP与降雨量偏相关系数空间分布')
plt.colorbar(im1, ax=ax1, label='偏相关系数')

# 绘制气温偏相关系数图
im2 = ax2.imshow(partial_corr_temperature, 
                 cmap='RdBu_r',
                 vmin=-1, 
                 vmax=1)
ax2.set_title('NPP与气温偏相关系数空间分布')
plt.colorbar(im2, ax=ax2, label='偏相关系数')

# 调整布局
plt.tight_layout()

# 保存图片
plt.savefig(os.path.join(output_path, 'partial_correlation_visualization.png'), 
            dpi=300, 
            bbox_inches='tight')

# 显示图形
plt.show()